package com.whz.gateway.common.web;

import com.whz.gateway.common.constants.GatewayConstant;
import com.whz.gateway.common.convert.Convert;
import com.whz.generic.spring.SpringUtils;
import lombok.experimental.UtilityClass;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.factory.rewrite.RewriteFunction;
import org.springframework.cloud.gateway.support.BodyInserterContext;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.ReactiveHttpOutputMessage;
import org.springframework.http.codec.CodecConfigurer;
import org.springframework.http.codec.HttpMessageReader;
import org.springframework.http.codec.ServerCodecConfigurer;
import org.springframework.http.codec.multipart.Part;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.util.MultiValueMap;
import org.springframework.web.reactive.function.BodyExtractors;
import org.springframework.web.reactive.function.BodyInserter;
import org.springframework.web.reactive.function.BodyInserters;
import org.springframework.web.reactive.function.client.ClientResponse;
import org.springframework.web.reactive.function.server.HandlerStrategies;
import org.springframework.web.reactive.function.server.ServerRequest;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Objects;
import java.util.function.Predicate;
import java.util.stream.Stream;

/**
 * webflux request body utils
 *
 * @author whz
 */
@Slf4j
@UtilityClass
public class WebBodyUtils {

    /**
     * 按照{@link Convert}的配置转换request body
     *
     * @param exchange exchange
     * @param convert  转换器配置
     * @param <T>      原始body类型
     * @param <R>      新body类型
     * @return request 已经改变的 ServerWebExchange
     */
    public static <T, R> Mono<Void> exchangeRequestBody(final GatewayFilterChain chain,
                                                        final ServerWebExchange exchange, final Convert<T, R> convert) {
        final Class<T> inClass = convert.getInClass();

        ServerRequest serverRequest = getCurrentServerRequest(exchange);

        RewriteFunction<T, R> rewriteFunction = convert.getRewriteFunction();

        Mono<R> modifiedBody = serverRequest.bodyToMono(inClass)
                                            .flatMap(originalBody -> Mono.defer(
                                                    () -> Mono.from(rewriteFunction.apply(exchange, originalBody))))
                                            .switchIfEmpty(
                                                    Mono.defer(() -> Mono.from(rewriteFunction.apply(exchange, null))));

        BodyInserter<Mono<R>, ReactiveHttpOutputMessage> bodyInserter =
                BodyInserters.fromPublisher(modifiedBody, convert.getOutClass());

        HttpHeaders headers = new HttpHeaders();
        headers.putAll(exchange.getRequest().getHeaders());
        headers.remove(HttpHeaders.CONTENT_LENGTH);
        if (convert.getContentType() != null) {
            headers.set(HttpHeaders.CONTENT_TYPE, convert.getContentType());
        }

        CachedHttpOutputMessage outputMessage = new CachedHttpOutputMessage(exchange, headers);
        return bodyInserter.insert(outputMessage, new BodyInserterContext())
                           .then(Mono.defer(() -> {
                               ServerHttpRequest decorator = decorate(exchange, headers, outputMessage);
                               return chain.filter(exchange.mutate().request(decorator).build());
                           }))
                           .onErrorResume(throwable -> {
                               log.error("body 转换异常.", throwable);
                               if (outputMessage.isCached()) {
                                   return outputMessage.getBody().map(DataBufferUtils::release)
                                                       .then(Mono.error(throwable));
                               }
                               return Mono.error(throwable);
                           });
    }

    ServerHttpRequestDecorator decorate(ServerWebExchange exchange, HttpHeaders headers,
                                        CachedHttpOutputMessage outputMessage) {
        return new LocalServerHttpRequestDecorator(exchange, headers, outputMessage);
    }


    /**
     * 重构当前response body,只处理APPLICATION_JSON_UTF8类型的数据
     *
     * @param exchange 上线文交换机
     * @param convert  response body修改器
     * @return 修改了response的上线文交换机
     * @see org.springframework.cloud.gateway.filter.factory.rewrite.ModifyResponseBodyGatewayFilterFactory
     */
    public static <T, R> ServerHttpResponse exchangeResponseBody(final ServerWebExchange exchange,
                                                                 Predicate<ServerWebExchange> predicate,
                                                                 Convert<T, R> convert) {
        // 重构response body
        return new LocalServerHttpResponseDecorator<>(predicate, exchange, convert);
    }

    /**
     * 更新path
     *
     * @param exchange exchange
     * @param newPath  新的path
     * @return ServerWebExchange
     */
    public static ServerWebExchange changePath(final ServerWebExchange exchange, final String newPath) {
        ServerHttpRequest newReq = exchange.getRequest().mutate().path(newPath).build();
        return exchange.mutate().request(newReq).build();
    }

    /**
     * @return HttpMessageReader List
     */
    public static List<HttpMessageReader<?>> getMessageReaders() {
        return SpringUtils.getBean(ServerCodecConfigurer.class).map(CodecConfigurer::getReaders).orElse(getOther());
    }

    private static List<HttpMessageReader<?>> getOther() {
        return HandlerStrategies.builder()
                                .codecs(config -> config.defaultCodecs().maxInMemorySize(GatewayConstant.MAX_BYTE_SIZE))
                                .build().messageReaders();
    }

    /**
     * get current request multiple part body
     *
     * @param exchange exchange
     * @return MMultiValueMap
     */
    public Mono<MultiValueMap<String, Part>> getCurrentMultiPartBody(final ServerWebExchange exchange) {
        return getCurrentServerRequest(exchange).body(BodyExtractors.toMultipartData());
    }

    /**
     * get current request wrapper
     *
     * @param exchange exchange
     * @return ServerRequest
     */
    private static ServerRequest getCurrentServerRequest(ServerWebExchange exchange) {
        return ServerRequest.create(exchange, getMessageReaders());
    }

    public ClientResponse getCurrentClientResponse(final ServerWebExchange exchange,
                                                   Publisher<? extends DataBuffer> body) {
        return ClientResponse.create(Objects.requireNonNull(exchange.getResponse().getStatusCode()),
                                     getMessageReaders())
                             .headers(headers -> headers.putAll(exchange.getResponse().getHeaders()))
                             .body(Flux.from(body)).build();
    }

    public boolean responseBodyText(final ServerWebExchange exchange) {
        MediaType contentType = exchange.getResponse().getHeaders().getContentType();
        return Stream.of(MediaType.TEXT_HTML).anyMatch(m -> m.isCompatibleWith(contentType));
    }

    public boolean responseBodyJson(final ServerWebExchange exchange) {
        MediaType contentType = exchange.getResponse().getHeaders().getContentType();
        return Stream.of(MediaType.APPLICATION_JSON).anyMatch(m -> m.isCompatibleWith(contentType));
    }

    public boolean requestBodyJson(final ServerWebExchange exchange) {
        MediaType contentType = exchange.getRequest().getHeaders().getContentType();
        return Stream.of(MediaType.APPLICATION_JSON).anyMatch(m -> m.isCompatibleWith(contentType));
    }
}